package com.pk.flink.apps.partitionAssigner;

import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.time.Duration;
import java.util.Arrays;
import java.util.Collection;
import java.util.Properties;

/**
 * RangeAssignor 分区分配规则测试用例
 * 涵盖：
 * 1. 基础Range分配逻辑（按主题分区、消费者字典序分配，余数归前N个消费者）
 * 2. 滚动重启后member.id变化导致分配错乱
 * 3. 静态成员（group.instance.id）保证分配粘性
 */
public class KafkaRangeAssignorConsumerApp {

    // Kafka集群地址
    private static String SERVERS = "master:9092,slave1:9092,slave2:9092";
    // 测试主题（需提前创建t0、t1两个主题，各3个分区）
    private static String TOPIC_T0 = "RangeAssignor_t0";
    private static String TOPIC_T1 = "RangeAssignor_t1";
    // 消费者组ID（所有测试用例共用同一组，模拟集群内消费者行为）
    private static String GROUP_ID = "range-assignor-test-group";

    private KafkaConsumer<String, String> consumer = null;

    /**
     * 基础配置初始化（通用配置）
     */
    @Before
    public void setUp() {
        Properties props = new Properties();
        // Kafka集群地址
        props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, SERVERS);
        // 消费者组ID
        props.put(ConsumerConfig.GROUP_ID_CONFIG, GROUP_ID);
        // 序列化配置
        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        // 显式指定分区分配器为RangeAssignor（Kafka默认也是，但显式声明更清晰）
        props.put(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, "org.apache.kafka.clients.consumer.RangeAssignor");
        // 自动提交offset（测试场景简化配置）
        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
        props.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, "1000");

        consumer = new KafkaConsumer<>(props);
    }

    /**
     * 测试用例1：基础RangeAssignor分配逻辑
     * 模拟2个消费者（C0、C1）订阅t0、t1（各3分区），验证分配结果：
     * C0: [t0p0, t0p1, t1p0, t1p1]
     * C1: [t0p2, t1p2]
     * （需启动两个该用例实例，分别模拟C0、C1）
     */
    @Test
    public void testRangeBasicAssign1() {
        // 订阅t0、t1两个主题
        consumer.subscribe(Arrays.asList(TOPIC_T0, TOPIC_T1), new ConsumerRebalanceListener() {
            @Override
            public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
                System.out.println("[Basic Range] 消费者撤销分区: " + partitions);
            }

            @Override
            public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
                System.out.println("[Basic Range] 消费者分配到分区: " + partitions);
                // 此处可打印当前消费者的member.id（需debug或通过KafkaAdminClient查询）
                System.out.println("[Basic Range] 当前消费者member.id: " + consumer.groupMetadata().memberId());
            }
        });

        // 持续拉取消息（模拟消费者运行）
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(java.time.Duration.ofMillis(Long.MAX_VALUE));
            for (ConsumerRecord<String, String> record : records) {
                System.out.printf("[Basic Range] key:%s, value:%s, topic:%s, partition:%d, offset:%d%n",
                        record.key(), record.value(), record.topic(), record.partition(), record.offset());
            }
        }
    }

    /**
     * 测试用例1：基础RangeAssignor分配逻辑
     * 模拟2个消费者（C0、C1）订阅t0、t1（各3分区），验证分配结果：
     * C0: [t0p0, t0p1, t1p0, t1p1]
     * C1: [t0p2, t1p2]
     * （需启动两个该用例实例，分别模拟C0、C1）
     */
    @Test
    public void testRangeBasicAssign2() {
        // 订阅t0、t1两个主题
        consumer.subscribe(Arrays.asList(TOPIC_T0, TOPIC_T1), new ConsumerRebalanceListener() {
            @Override
            public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
                System.out.println("[Basic Range] 消费者撤销分区: " + partitions);
            }

            @Override
            public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
                System.out.println("[Basic Range] 消费者分配到分区: " + partitions);
                // 此处可打印当前消费者的member.id（需debug或通过KafkaAdminClient查询）
                System.out.println("[Basic Range] 当前消费者member.id: " + consumer.groupMetadata().memberId());
            }
        });

        // 持续拉取消息（模拟消费者运行）
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(java.time.Duration.ofMillis(Long.MAX_VALUE));
            for (ConsumerRecord<String, String> record : records) {
                System.out.printf("[Basic Range] key:%s, value:%s, topic:%s, partition:%d, offset:%d%n",
                        record.key(), record.value(), record.topic(), record.partition(), record.offset());
            }
        }
    }


    /**
     * 测试用例2：滚动重启后member.id变化导致分配错乱
     * 模拟场景：
     * 1. 初始C0（member.id=C0）分配[t0p0,t0p1,t1p0,t1p1]，C1（member.id=C1）分配[t0p2,t1p2]
     * 2. 滚动重启后，C0变为member.id=C3，C1变为member.id=C2
     * 3. 因member.id字典序变化（C2 < C3），分配结果错乱：
     *    C3(原C0): [t0p2, t1p2]
     *    C2(原C1): [t0p0,t0p1,t1p0,t1p1]
     * （需先启动testRangeBasicAssign，记录member.id，再重启该用例模拟member.id变更）
     */
    @Test
    public void testRangeRebootAssignDisorder() {
        // 订阅t0、t1两个主题（与基础用例一致）
        consumer.subscribe(Arrays.asList(TOPIC_T0, TOPIC_T1), new ConsumerRebalanceListener() {
            @Override
            public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
                System.out.println("[Reboot Range] 消费者撤销分区: " + partitions);
            }

            @Override
            public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
                System.out.println("[Reboot Range] 滚动重启后分配到分区: " + partitions);
                System.out.println("[Reboot Range] 重启后消费者member.id: " + consumer.groupMetadata().memberId());
            }
        });

        // 持续拉取消息
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(java.time.Duration.ofMillis(Long.MAX_VALUE));
            for (ConsumerRecord<String, String> record : records) {
                System.out.printf("[Reboot Range] key:%s, value:%s, topic:%s, partition:%d, offset:%d%n",
                        record.key(), record.value(), record.topic(), record.partition(), record.offset());
            }
        }
    }

    /**
     * 测试用例3：静态成员（group.instance.id）保证分配粘性
     * 模拟场景：
     * 1. 给消费者设置固定的group.instance.id（I0、I1）
     * 2. 滚动重启后，即使member.id变化，分配结果仍保持：
     *    I0: [t0p0,t0p1,t1p0,t1p1]
     *    I1: [t0p2, t1p2]
     */
    @Test
    public void testRangeStaticMemberStickyAssign() {
        // 重新构建配置，添加静态成员ID（需为每个消费者设置唯一的instance.id）
        Properties staticProps = new Properties();
        staticProps.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, SERVERS);
        staticProps.put(ConsumerConfig.GROUP_ID_CONFIG, GROUP_ID);
        staticProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        staticProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        staticProps.put(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, "org.apache.kafka.clients.consumer.RangeAssignor");
        staticProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "true");
        staticProps.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, "1000");

        // 模拟消费者I0（需为另一个消费者实例设置为I1）
        // 注意：实际测试时需启动两个实例，分别设置I0和I1
        staticProps.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG, "I0");

        // 重建消费者（应用静态成员配置）
        consumer = new KafkaConsumer<>(staticProps);

        consumer.subscribe(Arrays.asList(TOPIC_T0, TOPIC_T1), new ConsumerRebalanceListener() {
            @Override
            public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
                System.out.println("[Static Range] 静态消费者撤销分区: " + partitions);
            }

            @Override
            public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
                System.out.println("[Static Range] 静态消费者分配到分区: " + partitions);
                System.out.println("[Static Range] 静态消费者instance.id: " + staticProps.getProperty(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG));
                System.out.println("[Static Range] 当前member.id: " + consumer.groupMetadata().memberId());
            }
        });

        // 持续拉取消息
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(Long.MAX_VALUE));
            for (ConsumerRecord<String, String> record : records) {
                System.out.printf("[Static Range] key:%s, value:%s, topic:%s, partition:%d, offset:%d%n",
                        record.key(), record.value(), record.topic(), record.partition(), record.offset());
            }
        }
    }

    /**
     * 资源释放
     */
    @After
    public void tearDown() {
        if (consumer != null) {
            consumer.close();
        }
    }
}
